import type { JSONSchemaType } from "ajv";
import {
  BaseTool,
  ChatMessage,
  JSONValue,
  Settings,
  ToolMetadata,
} from "llamaindex";

// prompt based on https://github.com/e2b-dev/ai-artifacts
const CODE_GENERATION_PROMPT = `You are a skilled software engineer. You do not make mistakes. Generate an artifact. You can install additional dependencies. You can use one of the following templates:\n

1. code-interpreter-multilang: "Runs code as a Jupyter notebook cell. Strong data analysis angle. Can use complex visualisation to explain results.". File: script.py. Dependencies installed: python, jupyter, numpy, pandas, matplotlib, seaborn, plotly. Port: none.

2. nextjs-developer: "A Next.js 13+ app that reloads automatically. Using the pages router.". File: pages/index.tsx. Dependencies installed: nextjs@14.2.5, typescript, @types/node, @types/react, @types/react-dom, postcss, tailwindcss, shadcn. Port: 3000.

3. vue-developer: "A Vue.js 3+ app that reloads automatically. Only when asked specifically for a Vue app.". File: app.vue. Dependencies installed: vue@latest, nuxt@3.13.0, tailwindcss. Port: 3000.

4. streamlit-developer: "A streamlit app that reloads automatically.". File: app.py. Dependencies installed: streamlit, pandas, numpy, matplotlib, request, seaborn, plotly. Port: 8501.

5. gradio-developer: "A gradio app. Gradio Blocks/Interface should be called demo.". File: app.py. Dependencies installed: gradio, pandas, numpy, matplotlib, request, seaborn, plotly. Port: 7860.

Provide detail information about the artifact you're about to generate in the following JSON format with the following keys:
  
commentary: Describe what you're about to do and the steps you want to take for generating the artifact in great detail.
template: Name of the template used to generate the artifact.
title: Short title of the artifact. Max 3 words.
description: Short description of the artifact. Max 1 sentence.
additional_dependencies: Additional dependencies required by the artifact. Do not include dependencies that are already included in the template.
has_additional_dependencies: Detect if additional dependencies that are not included in the template are required by the artifact.
install_dependencies_command: Command to install additional dependencies required by the artifact.
port: Port number used by the resulted artifact. Null when no ports are exposed.
file_path: Relative path to the file, including the file name.
code: Code generated by the artifact. Only runnable code is allowed.

Make sure to use the correct syntax for the programming language you're using. Make sure to generate only one code file. If you need to use CSS, make sure to include the CSS in the code file using Tailwind CSS syntax.
`;

// detail information to execute code
export type CodeArtifact = {
  commentary: string;
  template: string;
  title: string;
  description: string;
  additional_dependencies: string[];
  has_additional_dependencies: boolean;
  install_dependencies_command: string;
  port: number | null;
  file_path: string;
  code: string;
  files?: string[];
};

export type CodeGeneratorParameter = {
  requirement: string;
  oldCode?: string;
  sandboxFiles?: string[];
};

export type CodeGeneratorToolParams = {
  metadata?: ToolMetadata<JSONSchemaType<CodeGeneratorParameter>>;
};

const DEFAULT_META_DATA: ToolMetadata<JSONSchemaType<CodeGeneratorParameter>> =
  {
    name: "artifact",
    description: `Generate a code artifact based on the input. Don't call this tool if the user has not asked for code generation. E.g. if the user asks to write a description or specification, don't call this tool.`,
    parameters: {
      type: "object",
      properties: {
        requirement: {
          type: "string",
          description: "The description of the application you want to build.",
        },
        oldCode: {
          type: "string",
          description: "The existing code to be modified",
          nullable: true,
        },
        sandboxFiles: {
          type: "array",
          description:
            "A list of sandbox file paths. Include these files if the code requires them.",
          items: {
            type: "string",
          },
          nullable: true,
        },
      },
      required: ["requirement"],
    },
  };

export class CodeGeneratorTool implements BaseTool<CodeGeneratorParameter> {
  metadata: ToolMetadata<JSONSchemaType<CodeGeneratorParameter>>;

  constructor(params?: CodeGeneratorToolParams) {
    this.metadata = params?.metadata || DEFAULT_META_DATA;
  }

  async call(input: CodeGeneratorParameter) {
    try {
      const artifact = await this.generateArtifact(
        input.requirement,
        input.oldCode,
      );
      if (input.sandboxFiles) {
        artifact.files = input.sandboxFiles;
      }
      return artifact as JSONValue;
    } catch (error) {
      return { isError: true };
    }
  }

  // Generate artifact (code, environment, dependencies, etc.)
  async generateArtifact(
    query: string,
    oldCode?: string,
  ): Promise<CodeArtifact> {
    const userMessage = `
    ${query}
    ${oldCode ? `The existing code is: \n\`\`\`${oldCode}\`\`\`` : ""}
    `;
    const messages: ChatMessage[] = [
      { role: "system", content: CODE_GENERATION_PROMPT },
      { role: "user", content: userMessage },
    ];
    try {
      const response = await Settings.llm.chat({ messages });
      const content = response.message.content.toString();
      const jsonContent = content
        .replace(/^```json\s*|\s*```$/g, "")
        .replace(/^`+|`+$/g, "")
        .trim();
      const artifact = JSON.parse(jsonContent) as CodeArtifact;
      return artifact;
    } catch (error) {
      console.log("Failed to generate artifact", error);
      throw error;
    }
  }
}
